import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import Delaunay
from mesh_generator import export_mesh

def calculate_element_quality(vertices):
    """
    计算单个三角形单元的质量
    返回值在[0,1]之间，1表示等边三角形
    """
    # 计算边长
    edges = np.array([
        vertices[1] - vertices[0],
        vertices[2] - vertices[1],
        vertices[0] - vertices[2]
    ])
    edge_lengths = np.sqrt(np.sum(edges**2, axis=1))
    
    # 计算面积
    a, b, c = edge_lengths
    s = (a + b + c) / 2  # 半周长
    area = np.sqrt(s*(s-a)*(s-b)*(s-c))
    
    # 计算质量因子 (4*sqrt(3)*area)/(sum(edge_lengths^2))
    quality = 4 * np.sqrt(3) * area / np.sum(edge_lengths**2)
    
    return quality, edge_lengths

def check_mesh_quality(nodes, elements, well_pos, well_radius, n_trans):
    """检查网格质量"""
    qualities = []
    sizes = []
    dist_to_well = []
    
    for element in elements:
        vertices = nodes[element]
        
        # 计算单元质量
        quality, edge_lengths = calculate_element_quality(vertices)
        qualities.append(quality)
        
        # 计算单元尺寸（最长边）
        sizes.append(np.max(edge_lengths))
        
        # 计算到井的距离（单元中心到井的距离）
        center = np.mean(vertices, axis=0)
        dist = np.linalg.norm(center - well_pos)
        dist_to_well.append(dist)
    
    qualities = np.array(qualities)
    sizes = np.array(sizes)
    dist_to_well = np.array(dist_to_well)
    
    # 输出统计信息
    print("\n网格质量统计:")
    print(f"质量范围: [{qualities.min():.3f}, {qualities.max():.3f}]")
    print(f"平均质量: {qualities.mean():.3f}")
    print(f"标准差: {qualities.std():.3f}")
    
    print("\n网格尺寸统计:")
    print(f"尺寸范围: [{sizes.min():.2f}, {sizes.max():.2f}]")
    print(f"平均尺寸: {sizes.mean():.2f}")
    
    # 绘制质量分布图
    plt.figure(figsize=(15, 5))
    
    # 质量分布直方图
    plt.subplot(131)
    plt.hist(qualities, bins=50)
    plt.title('Element Quality Distribution')
    plt.xlabel('Quality')
    plt.ylabel('Count')
    
    # 尺寸与距离关系散点图
    plt.subplot(132)
    plt.scatter(dist_to_well, sizes, alpha=0.1)
    # 添加过渡层半径标记
    for i in range(n_trans + 1):
        r = well_radius * (2 ** i)
        plt.axvline(x=r, color='r', linestyle='--', alpha=0.5)
    plt.title('Element Size vs Distance to Well')
    plt.xlabel('Distance to Well (m)')
    plt.ylabel('Element Size (m)')
    
    # 质量空间分布
    plt.subplot(133)
    centers = np.array([np.mean(nodes[element], axis=0) for element in elements])
    plt.tripcolor(centers[:, 0], centers[:, 1], qualities, 
                 shading='gouraud', cmap='viridis')
    plt.colorbar(label='Quality')
    plt.title('Quality Distribution')
    plt.axis('equal')
    
    plt.tight_layout()
    plt.show()
    
    return qualities, sizes

def generate_refined_mesh(size=100.0, h_base=5.0, well_pos=(0.0, 0.0), well_radius=10.0, h_well=0.5, n_trans=3):
    """
    生成井周围多层加密的高质量三角形网格
    参数:
        size: 正方形边长的一半
        h_base: 基础网格尺寸
        well_pos: 井的位置 (x, y)
        well_radius: 最内层加密区域半径
        h_well: 井附近的最小网格尺寸
        n_trans: 过渡层数量
    """
    # 生成基础网格点
    n_base = int(2 * size / h_base) + 1
    x_base = np.linspace(-size, size, n_base)
    y_base = np.linspace(-size, size, n_base)
    
    # 计算各层的半径和网格尺寸
    radii = []    # 各层半径
    sizes = []    # 各层网格尺寸
    for i in range(n_trans + 1):
        # 半径按指数增长
        r = well_radius * (2 ** i)
        radii.append(r)
        # 网格尺寸按对数过渡
        h = h_well * (h_base/h_well) ** (i/n_trans)
        sizes.append(h)
    
    # 生成各层网格点
    x_points = [x_base]
    y_points = [y_base]
    
    for i in range(n_trans + 1):
        n_layer = int(2 * radii[i] / sizes[i]) + 1
        x_layer = np.linspace(well_pos[0]-radii[i], well_pos[0]+radii[i], n_layer)
        y_layer = np.linspace(well_pos[1]-radii[i], well_pos[1]+radii[i], n_layer)
        x_points.append(x_layer)
        y_points.append(y_layer)
    
    # 合并所有网格点
    x_all = np.unique(np.concatenate(x_points))
    y_all = np.unique(np.concatenate(y_points))
    X, Y = np.meshgrid(x_all, y_all)
    nodes = np.column_stack((X.flatten(), Y.flatten()))
    
    # 使用Delaunay三角剖分
    tri = Delaunay(nodes)
    elements = tri.simplices
    
    # 标记边界节点
    boundary_nodes = []
    tol = h_base/10
    
    for i, (x, y) in enumerate(nodes):
        if abs(x + size) < tol:  # 左边界
            boundary_nodes.append((i+1, 1))
        elif abs(x - size) < tol:  # 右边界
            boundary_nodes.append((i+1, 2))
        elif abs(y + size) < tol:  # 下边界
            boundary_nodes.append((i+1, 3))
        elif abs(y - size) < tol:  # 上边界
            boundary_nodes.append((i+1, 4))
    
    # 检查网格质量
    qualities, sizes = check_mesh_quality(nodes, elements, well_pos, well_radius, n_trans)
    
    # 输出质量警告
    if np.min(qualities) < 0.3:
        print("\n警告: 存在低质量单元!")
        poor_elements = np.where(qualities < 0.3)[0]
        print(f"低质量单元数量: {len(poor_elements)}")
        print(f"低质量单元位置: {poor_elements}")
    
    # 输出网格信息
    print("\n网格信息:")
    print(f"基础网格尺寸: {h_base:.1f} m")
    print(f"井附近网格尺寸: {h_well:.1f} m")
    print(f"过渡层数量: {n_trans}")
    for i in range(n_trans + 1):
        print(f"第{i}层: 半径={radii[i]:.1f}m, 网格尺寸={sizes[i]:.2f}m")
    print(f"节点总数: {len(nodes)}")
    print(f"单元总数: {len(elements)}")
    print(f"边界节点数: {len(boundary_nodes)}")
    
    return nodes, elements, boundary_nodes

def plot_refined_mesh(nodes, elements, well_pos, well_radius, n_trans=3):
    """绘制加密网格"""
    plt.figure(figsize=(10, 10))
    plt.triplot(nodes[:, 0], nodes[:, 1], elements, 'b-', linewidth=0.5)
    plt.plot(nodes[:, 0], nodes[:, 1], 'r.', markersize=1)
    
    # 绘制各层加密区域
    for i in range(n_trans + 1):
        r = well_radius * (2 ** i)
        circle = plt.Circle(well_pos, r, fill=False, color='red', 
                          linestyle='--', alpha=0.5)
        plt.gca().add_patch(circle)
    
    plt.plot(well_pos[0], well_pos[1], 'r*', markersize=10, label='Well')
    
    plt.axis('equal')
    plt.grid(True)
    plt.title('Refined Triangle Mesh')
    plt.legend()
    plt.show()

if __name__ == "__main__":
    # 网格参数
    size = 50.0          # 区域大小
    h_base = 4.0         # 基础网格尺寸
    well_pos = (0.0, 0.0) # 井的位置
    well_radius = 5.0    # 最内层加密区域半径
    h_well = 0.5         # 井附近的最小网格尺寸
    n_trans = 3          # 过渡层数量
    
    # 生成网格
    nodes, elements, boundary_nodes = generate_refined_mesh(
        size=size,
        h_base=h_base,
        well_pos=well_pos,
        well_radius=well_radius,
        h_well=h_well,
        n_trans=n_trans
    )
    
    # 绘制和导出
    plot_refined_mesh(nodes, elements, well_pos, well_radius, n_trans)
    export_mesh(nodes, elements, boundary_nodes, 'square_mesh_refined') 